package handlers

import (
	"bufio"
	"bytes"
	"fmt"
	"generalhandler/conf"
	"generalhandler/utils"
	"github.com/bitly/go-simplejson"
	"io"
	"os"
	"time"
)

type CrawlerHandler struct {
}

func (ch *CrawlerHandler) Handle(bytesMsg []byte, log Logger, pub Publisher) error {
	start := time.Now().Unix()
	msg, err := simplejson.NewJson(bytesMsg)
	if nil != err {
		log.Warn("failed to NewJson %s, cuz %s", string(bytesMsg), err.Error())
		return nil
	}
	strDomain, err := msg.Get("domain").String()
	if nil != err {
		log.Warn("found no domain in %s", string(bytesMsg))
		return nil
	}
	_, err = msg.Get("tkid").String()
	if nil != err {
		log.Warn("found no tkid in %s", string(bytesMsg))
		return nil
	}
	strSubDomain, err := msg.Get("subdomain").String()
	if nil != err {
		log.Warn("found no subdomain in %s", string(bytesMsg))
		return nil
	}
	strPort, err := msg.Get("port").String()
	if nil != err {
		log.Warn("found no port in %s", string(bytesMsg))
		return nil
	}
	strPrefix, err := msg.Get("prefix").String()
	if nil != err {
		log.Warn("found no prefix in %v", string(bytesMsg))
		return nil
	}
	//lzy add user agent
	strUserAgent, err := msg.Get("useragent").String()
	if nil != err {
		log.Warn("found no useragent in %v", string(bytesMsg))
		return nil
	}
	//var userAgent map[string]string
	//userAgent = make(map[string]string)
	//fmt.Println(userAgent)
	//userAgent["user-agent"] = strUserAgent

	strIgnore, _ := msg.Get("ignore").String()
	var ignoreMap map[string]interface{}
	ignoreMap = make(map[string]interface{})
	utils.StringToMap(strIgnore, ignoreMap)

	bin := "/usr/local/arachni/bin/arachni"
	var args []string
	//lzy add
	ua := fmt.Sprintf("--http-user-agent=%s", strUserAgent)
	args=append(args,ua)

	args = append(args, "--checks=-")
	args = append(args, "--report-save-path=/dev/null")
	out := fmt.Sprintf("%s/out", utils.GetSelfBinDirPath())
	defer os.Remove(out)
	plugin := fmt.Sprintf("--plugin=urls_dump:dump_path=%s", out)
	args = append(args, plugin)
	args = append(args, "--browser-cluster-pool-size=3")

	proxy := conf.CrawlerProxyConfigInstance().Proxy
	auth := conf.CrawlerProxyConfigInstance().Auth
	t := conf.CrawlerProxyConfigInstance().Type
	if 0 != len(proxy) {
		args = append(args, fmt.Sprintf("--http-proxy=%s", proxy))
		if 0 != len(auth) {
			args = append(args, fmt.Sprintf("--http-proxy-authentication=%s", auth))
		}
		if 0 != len(t) {
			args = append(args, fmt.Sprintf("--http-proxy-type=%s", t))
		}
	}

	url := fmt.Sprintf("%s://%s:%s", strPrefix, strSubDomain, strPort)
	args = append(args, url)

	scanresult := make(map[string]interface{})
	buf := new(bytes.Buffer)
	std := new(bytes.Buffer)
	cmdLog := utils.NewCmdLogInstance(utils.GetSelfBinDirPath(), 10*time.Minute)
	err = utils.ExecCmd(10*60*60, std, buf, cmdLog, bin, args...)
	if nil != err {
		log.Error("failed to crawler %s, cuz %s-%s", strDomain, err.Error(), buf.String())
		scanresult["scan_error"] = buf.String()
	}
	//fmt.Println(std.String())
	scanresult["scan_type"] = "crawler"
	msg.Set("start", start)
	msg.Set("end", time.Now().Unix())

	f, err := os.Open(out)
	if nil != err {
		log.Error("failed to open %s, cuz %s", out, err.Error())
		return nil
	}
	defer f.Close()
	rdr := bufio.NewReader(f)
	var urls []string
	for ; ; {
		line, _, err := rdr.ReadLine()
		if io.EOF == err {
			break
		}
		if nil != err {
			break
		}
		urls = append(urls, string(line))
		if 50 == len(urls) {
			scanresult["urls"] = urls
			msg.Set("scan_result", scanresult)
			urls = nil
			newMsg, err := msg.MarshalJSON()
			if nil != err {
				log.Error("failed to marshal %v, cuz %s", msg, err.Error())
				continue
			}
			for _, handle := range conf.MQConfigInstance().NextHandlers {
				if _, ok := ignoreMap[handle]; ok {
					continue
				}
				pub.PublishMessage(handle, newMsg)
			}
		}
	}
	if 0 != len(urls) {
		scanresult["urls"] = urls
		msg.Set("scan_result", scanresult)
		urls = nil
		newMsg, err := msg.MarshalJSON()
		if nil != err {
			log.Error("failed to marshal %v, cuz %s", msg, err.Error())
			return nil
		}
		for _, handle := range conf.MQConfigInstance().NextHandlers {
			if _, ok := ignoreMap[handle]; ok {
				continue
			}
			pub.PublishMessage(handle, newMsg)
		}
	}

	return nil
}
